# -*- coding: utf-8 -*-
"""
PBAR_HighZFinderPerformance.py

- Determines the performance of the merged-merged marker.
- performance values for each combination of settings saved to file

Created on Thu May 15 12:02:13 2014

@author: jkwong
"""

import glob, os, gc, objgraph, cPickle, time
import numpy as np
import PBAR_Zspec, PBAR_Cargo
import matplotlib.pyplot as plt
from matplotlib import mpl
from scipy import ndimage
reload(PBAR_Zspec)

dataPath = r'E:\PBAR\data4\BasicScansStandardWidthPickle'
data0Path = r'E:\PBAR\data4\HighZFinder'
plotSaveDir = r'E:\PBAR\data4\HighZFinder'

num = 11

figureSize = (16, 10)

# Global parameters
MINBIN = 1
newWidth = 600
newHeight = 136
xCargoInternalBounds = np.array((38, 550))
yCargoInternalBounds = np.array((13, 106))

#cargoCountRange = np.array([0, 4e7]) # only consider parts of the zspec image that have rad below this value
cargoCountRange = np.array([0, 3e7]) # only consider parts of the zspec image that have rad below this value

ZspecCountRange = np.array([0, 275])

# make different size windows
# x and y
windowList = []
windowList.append(np.array((3, 1)))
windowList.append(np.array((5, 1)))
windowList.append(np.array((7, 1)))
windowList.append(np.array((4, 2)))
windowList.append(np.array((7, 2)))
# thresholds
discrimThresholdList = np.array((0, 2))
# fractions
fractionPixelsList = np.array((0.8, 1.0))

#tolerance = [3, 1]

# get linear discriminant coefficients, polynomial fits
datasetDescription = PBAR_Zspec.ReadCargoDataDescriptionFile(r'E:\PBAR\data4\CargoSet2.txt')
for datasetIndex, datasetDescript in enumerate(datasetDescription['scanID']):
    print(datasetIndex)

#    if datasetIndex < 65:
#        continue
#    if datasetIndex < 61:
#        continue
#    if datasetIndex > 64:
#        continue
#    if datasetIndex == 63:
#        continue
    
    if datasetIndex < 66:
        continue
    
    zspecScanNumber = np.int(datasetDescript)

    # Read Zspec
    filenameZspec = '%s-FDFC-All_SW.dat' %datasetDescription['scanID'][datasetIndex]
    fullFilenameZspec = os.path.join(dataPath, filenameZspec)
    
    # Read Cargo
    filenameCargo = 'PBAR-%s.cargoimageSW.dat' %datasetDescription['dataFile'][datasetIndex]
    fullFilenameCargo = os.path.join(dataPath, filenameCargo)
    
    # read in the cargo image
    print('Loading %s' %fullFilenameCargo)
    with open(fullFilenameCargo, 'rb') as fid:
        datCargoStandardWidth = cPickle.load(fid)    
    datCargoSmall = PBAR_Zspec.CargoReduceSize(datCargoStandardWidth, newHeight, newWidth)
    
    energy, datZspecStandardWidth = PBAR_Zspec.ReadZspecBasicScanPickle(fullFilenameZspec)
    datZspecSmall = PBAR_Zspec.ZspecBasicReduceSize(datZspecStandardWidth, energy, newWidth)
    # calculate the features for reduced size zspec image
    discrimSmall = PBAR_Zspec.CalculateFeaturesZspecBasic(datZspecSmall, energy, MINBIN)    
    
    # Read in marker files
    filenameMarker = filenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    fullFilenameMarker = fullFilenameCargo.replace('cargoimageSW.dat', 'cargomarkerSW')
    
    # some don't have marker files
    if os.path.exists(fullFilenameMarker):
        markerStandardWidth0 = PBAR_Cargo.ReadCargoMarker(fullFilenameMarker)
    else:
        markerStandardWidth0 = []
    
    # Modify the marker files for truncated
    multiplier = newWidth/float(datZspecStandardWidth.shape[0])
    offset = 0.0
    markerStandardWidth = PBAR_Zspec.ModifyMarkersXPosition(markerStandardWidth0, multiplier, offset)

    numberMarkers = len(markerStandardWidth)
    
    # determine if marker is high z (DU or W) or not (Fe or Pb)
    highZMaterial = np.zeros(numberMarkers) == 1 # start with all false
    for markerIndex, marker in enumerate(markerStandardWidth):
        # s = special, w = tungsten, du = depleted uranium
        if ( (marker['target'][0].lower() == 's') or (marker['target'][0].lower() == 'w') or (marker['target'][0].lower() == 'd') or (marker['target'][0].lower() == 't')):
            highZMaterial[markerIndex] = True
            
    # make masks
    markMaskList = []
    for markerIndex, marker in enumerate(markerStandardWidth):
        markerXRange = np.arange(marker['rec_left'], marker['rec_right']+1).astype(int)
#        markerYRange = np.arange(marker['rec_top'], marker['rec_bottom']+1).astype(int)
        markerYRange = np.arange(min(marker['rec_top'], marker['rec_bottom']), max(marker['rec_top'], marker['rec_bottom'])+1).astype(int)
        
        maskTemp = np.zeros((newWidth, newHeight)) == 1
        maskTemp[markerXRange[0]:markerXRange[-1], markerYRange[0]:markerYRange[-1]] = True
        markMaskList.append(maskTemp)
    
    # loop through algorithm settings
    for windowIndex, window in enumerate(windowList):
        for discrimThresholdIndex, discrimThreshold in enumerate(discrimThresholdList):
            for fractionPixelsIndex, fractionPixels in enumerate(fractionPixelsList):
                tStart = time.time()                
                
                mergedMergedPerformance = []
                
                # makes them all have the same reference, I think
#                highZMarkersCrossRef = [{'mergedMergedIndex':[]}] * numberMarkers

                highZMarkersCrossRef = []
                for i in xrange(numberMarkers):
                    highZMarkersCrossRef.append({})
                    highZMarkersCrossRef[i]['mergedMergedIndex'] = []
                
                fullFilenamePotential = \
                    fullFilenameMarker.replace('cargomarkerSW', \
                        '%d_%d_%d_%3.2f.potential' %(window[0], window[1], discrimThreshold, fractionPixels) )
                print('Working on %s' %fullFilenamePotential)
                
                # load original high markers
                with open(fullFilenamePotential, 'rb') as fid:
                    print('Load Original Markers: %s' %fullFilenamePotential)
                    highZMarkersList = cPickle.load(fid)
                # load the merged markers
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMerged')
                with open(fullFilenameMarkerMerged, 'rb') as fid:
                    print('Load Merged Markers: %s' %fullFilenameMarkerMerged)
                    highZMergedMarkersList = cPickle.load(fid)
                # Save to file
                fullFilenameMarkerMerged = fullFilenamePotential.replace('potential', 'highZMergedMerged')
                with open(fullFilenameMarkerMerged, 'rb') as fid:
                    print('Load Merged Merged Markers: %s' %fullFilenameMarkerMerged)
                    highZMergedMergedMarkersList = cPickle.load(fid)

                print('time elapsed1: %3.3f' %(time.time() - tStart))

                # number of times each truth marker was marked (not necessarily correctly)
                markerMergedMergedOverlapCount = np.zeros(numberMarkers) # okay if len is zero.
                # if mergedmerged marker is a false positive
                # - if mm doesn't intersect with a truth marker
                # - if mm does intersect with truth marker but it is not high-z
                mergedMergedFalsePositives = np.zeros(len(highZMergedMergedMarkersList)) == 1
                # mm overlaps with a high-z truth marker
                mergedMergedTruePositives = np.zeros(len(highZMergedMergedMarkersList)) == 1
#                # if the material tagged was actually high-z
#                mergedMergedHighZTruth = np.zeros(len(highZMergedMergedMarkersList)) == 1
                # material name
                # material target or background
                mergedMergedMaterial = []
                
                # if none of the merged merged finds the high z marker
                markerFalseNegatives = np.zeros(numberMarkers) == 1
                markerPrediction = np.zeros(numberMarkers) == 1
                
                if highZMergedMergedMarkersList == []:
                    print('No high-z merged merged markers found. Skipping')

                else:                                
                    # make list in case some markers find more than one.
    
                    # loop over the merged-merged markers
                    for highZMergedMergedMarkersIndex,highZMergedMergedMarkers in enumerate(highZMergedMergedMarkersList):
                        print('Checking merged-merged %d/%d' %(highZMergedMergedMarkersIndex+1, len(highZMergedMergedMarkersList)))
                        mergedMergedPerformanceTemp = {}
                        mergedMergedPerformanceTemp['markerFoundIndex'] = [] # if overlap found store number of truth marker
                        mergedMergedPerformanceTemp['markerFoundOverlap'] = []  # store number of pixels overlap
    
                        # track how many times that merged merged marker intersects with an actual object
                        intersectionCount = 0
                        # loop over the truth markers
                        for markMaskIndex, markMask in enumerate(markMaskList):
#                            print('time elapsed2: %3.3f' %(time.time() - tStart))
                            intersectionMask = markMask & highZMergedMergedMarkers['mask']
                            intersectionOverlap = intersectionMask.flatten().sum()
#                            print('time elapsed3: %3.3f' %(time.time() - tStart))
                            # determine if there was overlap between merge-merge and truth marker
                            if (intersectionOverlap > 0):
                                print('MM %d intersects with marker %d'  %(highZMergedMergedMarkersIndex, markMaskIndex))
                                # track target name
                                mergedMergedMaterial.append(markerStandardWidth[markMaskIndex]['target'])
                                markerMergedMergedOverlapCount[markMaskIndex] += 1 # increment if intersects 
                                mergedMergedPerformanceTemp['markerFoundIndex'].append(markMaskIndex)
                                mergedMergedPerformanceTemp['markerFoundOverlap'].append(intersectionOverlap)
                                markerPrediction[markMaskIndex] = True
                                # store the index of the MM marker that found it
                                highZMarkersCrossRef[markMaskIndex]['mergedMergedIndex'].append(highZMergedMergedMarkersIndex)
#                                print('time elapsed3e: %3.3f' %(time.time() - tStart))
                                # if overlap and is high-z then it is a true positive
                                if (highZMaterial[markMaskIndex]):
                                    mergedMergedTruePositives[highZMergedMergedMarkersIndex] = True
#                                    mergedMergedHighZTruth[highZMergedMergedMarkersIndex] = True
                                else: # if overlaps but it is not high z
                                    mergedMergedFalsePositives[highZMergedMergedMarkersIndex] = True
                                intersectionCount += 1
#                                print('time elapsed3f: %3.3f' %(time.time() - tStart))
#                            print('time elapsed4: %3.3f' %(time.time() - tStart))
                        mergedMergedPerformance.append(mergedMergedPerformanceTemp)
                        # it's a false positve if it doesn't intersect with any high-z markers
                        if ( (intersectionCount == 0) & (sum(highZMaterial) > 0) ):
                            print('MM %d is a false positive'  %highZMergedMergedMarkersIndex)
                            mergedMergedFalsePositives[highZMergedMergedMarkersIndex] = True
                            # if the marker tagged background
                            mergedMergedMaterial.append('background')
                        print highZMarkersCrossRef
                            
                    # false negative - no intersections and is high Z
                    for markMaskIndex,markMask in enumerate(markMaskList):
                        # if nothing in cross reference and is high-z
                        if ((highZMarkersCrossRef[markMaskIndex]['mergedMergedIndex'] == []) & highZMaterial[markMaskIndex]):
                            markerFalseNegatives[markMaskIndex] == True                    

                ## Small Cargo image ##
                plotMarkers = True
                plotPotential = False
                plotMergedMerged = True
                
                figureSize = (23, 13)
                
                plt.figure(figsize = figureSize)
                plt.grid()
                
                plt.imshow(datCargoSmall.T, interpolation = 'nearest', aspect='auto', cmap = plt.cm.Greys_r)
                #plt.imshow(discrimSmall['discrim'].T, interpolation = 'nearest', aspect='auto', cmap = cm.Greys_r)
                
                plt.colorbar()
                plt.ylabel('Detector Number')
                plt.xlabel('Time Slice')
                
                plt.axis((0, 600, 136, 0))
                
                if plotMergedMerged:
                    for mIndex, m in enumerate(highZMergedMergedMarkersList):
                        maskmask = np.ma.masked_where(m['mask'] == 0, m['mask'])
                #                        plt.imshow(maskmask.T, cmap=mpl.cm.jet_r, interpolation='nearest', aspect='auto')
                        plt.imshow(maskmask.T, cmap=plt.cm.Oranges_r, interpolation='nearest', aspect='auto')
                    for mIndex, m in enumerate(highZMergedMergedMarkersList):
                        if (mergedMergedTruePositives[mIndex]):
                            boxColor = 'g' # if correct
                        else:
                            boxColor = 'r'  # if incorrectly labeled
                        plt.text(m['xMean']+5, m['yMean']-1, '%d' % (mIndex+1), color = boxColor, fontsize = 16, fontweight='bold', alpha = 0.75)                    
                    
                if plotMarkers:
                #                    marker = markerStandardWidth
                    for markIndex, mark in enumerate( markerStandardWidth):
                        
                        if ( markerPrediction[markIndex] == highZMaterial[markIndex]):
                            boxColor = 'g' # if correct
                        else:
                            boxColor = 'r'  # if incorrectly labeled
                            
                        x = mark['rec_left'], mark['rec_right'], mark['rec_right'], mark['rec_left'], mark['rec_left']
                        y = mark['rec_bottom'], mark['rec_bottom'], mark['rec_top'], mark['rec_top'], mark['rec_bottom']
                        plt.plot(x, y, boxColor)
                #        plt.scatter(mark['x'], mark['y'], marker = '.', s = 14, color = boxColor, linewidth = 7, alpha = 0.8)
                        plt.plot(mark['x'], mark['y'], marker = '.', color = boxColor, alpha = 0.8, markersize = 7)
                        
                        if 'left' in mark:
                            plt.plot(mark['left']['x'], mark['left']['y'], 'gx')
                        if 'right' in mark:
                            plt.plot(mark['right']['x'], mark['right']['y'], 'go')
                        plt.text(max(x), min(y), mark['target'], color = 'b', fontsize = 16, fontweight='bold', alpha = 0.6)
                
                if plotPotential:
                    for (potIndex, pot) in enumerate(highZMarkersList):
                #        plt.plot(pot['boxCenter'][0], pot['boxCenter'][1], 'xb', markersize = 14)
                        plt.scatter(pot['boxCenter'][0], pot['boxCenter'][1], color = 'm', \
                            s = 50 * pot['window'].mean(), linewidth = pot['window'].mean()/2.0, edgecolor = 'magenta', facecolor = 'magenta', marker = 'x', alpha = 0.5)
                
                plt.title('%s, Rad, window: %d, %d, disrim thresh: %d, fractionPixels %3.2f \n #Truth Markers, high-z: %d, %d, # MM %d, #FP predictions %d, #TP predictions %d, #TP %d, #FN %d' \
                    %(filenameZspec, window[0], window[1], discrimThreshold, fractionPixels, len(markerFalseNegatives), \
                    sum(highZMaterial), len(mergedMergedFalsePositives), sum(mergedMergedFalsePositives), sum(mergedMergedTruePositives), \
                    sum(   (markerPrediction == highZMaterial) & (highZMaterial == 1)   ), sum(markerFalseNegatives)))                
                
                # save image to file
                plotFullFilename = os.path.join(plotSaveDir, '%s_%d_%d_%d_%3.2f_Cargo_MM.png' %(zspecScanNumber, window[0], window[1], discrimThreshold, fractionPixels))
                plt.savefig(plotFullFilename)
                print('Wrote: %s' %plotFullFilename)
                # plt.show()
                plt.close('all')

                print('Number of truth markers: %d' %len(markerFalseNegatives))
                print('Number of high-z truth markers: %d' %sum(highZMaterial))
                print('Number of non high-z truth markers: %d' %sum(~highZMaterial))
                print('Number of false positives predictions: %d,  %d' %(sum(mergedMergedFalsePositives),len(mergedMergedFalsePositives)) )
                print('Number of true positive predictions: %d, %d' %(sum(mergedMergedTruePositives), len(mergedMergedTruePositives))   )
                print('Number of true positives found: %d, Number of actual positives: %d' %(sum(   (markerPrediction == highZMaterial) & (highZMaterial == 1)   ), sum(highZMaterial))  )
                print('Number of false negatives: %d / %d' %(sum(markerFalseNegatives), len(markerFalseNegatives)) )
                try:
                    success = sum( (markerMergedMergedOverlapCount > 0) == highZMaterial) / np.float(len(highZMaterial))
                except:
                    success = np.nan
                print('Success : %3.3f' %success)
                
                # save results to file
                dat = {}
                dat['mergedMergedPerformance'] = mergedMergedPerformance
                dat['highZMarkersCrossRef'] = highZMarkersCrossRef
                dat['markerMergedMergedOverlapCount'] = markerMergedMergedOverlapCount
                dat['mergedMergedFalsePositives'] = mergedMergedFalsePositives
                dat['mergedMergedTruePositives'] = mergedMergedTruePositives
                dat['markerPrediction'] = markerPrediction
                dat['markerFalseNegatives'] = markerFalseNegatives
                dat['highZMaterial'] = highZMaterial
                dat['success'] = success
                
                fullFilenamePerformance = fullFilenamePotential.replace('potential', 'performance')
                with open(fullFilenamePerformance, 'wb') as fid:
                    print('Write: %s' %fullFilenamePerformance)
                    cPickle.dump(dat, fid, 2)
                print('time elapsed3: %3.3f' %(time.time() - tStart))
#                break
#            break
#        break
#    break
    
